package com.huahua.bigdata.sparksql;


import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.expressions.Aggregator;

import java.util.*;

public class MyCityRemarkUDAF extends Aggregator<String, MyCityRemarkBuffer, String> {
    @Override
    public MyCityRemarkBuffer zero() {
        return new MyCityRemarkBuffer(0L, new HashMap<String, Long>());
    }

    @Override
    // TODO 将函数的输入值和缓冲区的数据进行聚合处理
    public MyCityRemarkBuffer reduce(MyCityRemarkBuffer buffer, String city) {
        buffer.setCount(buffer.getCount() + 1);
        Map<String, Long> cityMap = buffer.getCityMap();
        Long cityCount = cityMap.get(city);
        if ( cityCount == null ) {
            cityMap.put(city, 1L);
        } else {
            cityMap.put(city, cityCount + 1);
        }
        buffer.setCityMap(cityMap);
        return buffer;
    }

    @Override
    // 合并缓冲区
    public MyCityRemarkBuffer merge(MyCityRemarkBuffer b1, MyCityRemarkBuffer b2) {
        b1.setCount( b1.getCount() + b2.getCount() );
        Map<String, Long> cityMap1 = b1.getCityMap();
        Map<String, Long> cityMap2 = b2.getCityMap();
        /**
         * c1 : {beijing : 10, tianjing: 20, baoding : 30}
         * c2 : {beijing : 40, tianjing : 50, sjz: 60}
         * ------------------------------------------------
         * c3 : {beijing : 50, tianjin : 70, baoding : 30, sjz ： 60}
         * 1. 将map1保持不变
         * 2. 对map2进行遍历
         * 3. 如果map2 中的key 在map1 存在, 那么合并数据
         * 4. 如果map2 中的key 不在 map1 存在, 那么直接添加即可
         * */
        final Iterator<String> iterator = cityMap2.keySet().iterator();
        while (iterator.hasNext()) {
            String key = iterator.next();

            Long v1 = cityMap1.get(key);
            Long v2 = cityMap2.get(key);
            if ( v1 == null ){
                cityMap1.put(key, v2);
            } else {
                cityMap1.put(key, v1 + v2 );
            }
        }
        b1.setCityMap(cityMap1);
        return b1;
    }

    @Override
    public String finish(MyCityRemarkBuffer buffer) {
        StringBuffer ss = new StringBuffer();
        Long total = buffer.getCount();
        Map<String, Long> cityMap = buffer.getCityMap();
        List<CityCount> ccs = new ArrayList<CityCount>();
        cityMap.forEach(
                (k, v) -> {
                    ccs.add( new CityCount(k, v));
                }
        );
        // TODO 对List进行排序
        Collections.sort(ccs);

        CityCount cityCount0 = ccs.get(0);
        long pc0 = cityCount0.getCount() * 100 / total;
        ss.append(cityCount0.getCityName() + " "+ pc0 +"%" );

        CityCount cityCount1 = ccs.get(1);

        long pc1 = cityCount1.getCount() * 100 / total;
        ss.append(cityCount1.getCityName() + " "+ pc1 +"%" );

        if ( ccs.size() > 2 ) {
            ss.append("其他"+(100 - pc0 - pc1)+"%");
        }
        return null;
    }

    @Override
    public Encoder<MyCityRemarkBuffer> bufferEncoder() {
        return Encoders.bean(MyCityRemarkBuffer.class);
    }

    @Override
    public Encoder<String> outputEncoder() {
        return Encoders.STRING();
    }
}
